# -*- coding: utf-8 -*-
from typing import Iterable, Dict, Callable, Union
# import pendulum
import pendulum
from pytz import timezone
from airflow.exceptions import AirflowSkipException
from airflow.models import BaseOperator
from airflow.utils.log.logging_mixin import LoggingMixin
from utils.macros import cst_monthday, cst_weekday, date_add


class ConditionRunMixin(LoggingMixin):
    """
    使 Operator 仅在满足特定条件时运行，否则跳过任务

    使用方式1 混入原生Operator，并传入 should_run 参数：
    ```
    from airflow.operators.bash import BashOperator as AirflowBashOperator  # 引入原生 Operator
    from utils.operators.mixins.condition_run_mixin import ConditionRunMixin  # 引入 Mixin
    from utils.macros import macros_dict, cst_weekday, cst_monthday, date_add

    # 混合 Operator 与 Mixin，使 Operator 具有自动跳过执行的能力
    class BashOperator(ConditionRunMixin, AirflowBashOperator):
        # 无需实现任何方法
        pass

    BashOperator(
      task_id='run_only_w1',
      python_callable=func,
      # 根据 execution_date 判断日期，仅周一运行
      should_run=lambda ctx: cst_weekday(date_add(ctx['ti'].execution_date, 1)) == '1',
    )
    ```

    使用方式2 子类实现 should_run 方法，再将子 Mixin 类混入 Operator：
    ```
    from airflow.operators.bash import BashOperator  # 引入原生 Operator
    from utils.operators.mixins.condition_run_mixin import ConditionRunMixin  # 引入 Mixin
    from utils.macros import macros_dict, cst_weekday, cst_monthday, date_add

    # 混合 Operator 与 Mixin，使 Operator 具有自动跳过执行的能力
    class WeekDay1OnlyOperator(ConditionRunMixin, BashOperator):
        def should_run(self, ctx):
            return cst_weekday(date_add(ctx['ti'].execution_date, 1)) == '1'

    WeekDay1OnlyOperator(
      task_id='run_only_w1',
      python_callable=func,
    )
    ```
    """

    def __init__(self, should_run: Callable = None, *args, **kwargs):
        super().__init__(*args, **kwargs)
        if should_run is None and hasattr(self, 'should_run'):
            should_run = self.should_run
        if not callable(should_run):
            raise NotImplementedError()
        self._should_run = should_run

    def pre_execute(self, context: Dict = None):
        if not self._should_run(context):
            self.log.info('should run == False')
            raise AirflowSkipException

        # 获取最顶层的 Operator 类型，执行 pre_execute
        cls = next(c for c in self.__class__.__bases__ if issubclass(c, BaseOperator))
        cls.pre_execute(self, context)


class DayOfWeekOnlyMixin(ConditionRunMixin):
    """
    仅每周某几天运行
    :param days_of_week
    :type int or list[int]
    """
    def __init__(self, days_of_week: Union[int, Iterable[int]], *args, **kwargs):
        if not isinstance(days_of_week, Iterable):
            days_of_week = [days_of_week]

        if not all(0 <= day <= 7 for day in days_of_week):
            raise Exception('DayOfWeekOnlyMixin param ERROR: `days_of_week` should be provide and between [0, 6]')

        super(DayOfWeekOnlyMixin, self).__init__(*args, **kwargs)
        self.days_of_week = [str(0 if day == 7 else day) for day in days_of_week]

    def should_run(self, ctx):
        return cst_weekday(date_add(ctx['ti'].execution_date, 1)) in self.days_of_week


class DayOfMonthOnlyMixin(ConditionRunMixin):
    """
    仅每月某几天运行
    :param days_of_month
    :type int or list[int]
    """
    def __init__(self, days_of_month: Union[int, Iterable[int]], *args, **kwargs):
        if not isinstance(days_of_month, Iterable):
            days_of_month = [days_of_month]

        if not all(1 <= day <= 31 for day in days_of_month):
            raise Exception('DayOfMonthOnlyMixin param ERROR: `days_of_month` should be provide and between [1, 31]')

        super(DayOfMonthOnlyMixin, self).__init__(*args, **kwargs)
        self.days_of_month = [str(n).zfill(2) for n in days_of_month]

    def should_run(self, ctx):
        return cst_monthday(date_add(ctx['ti'].execution_date, 1)) in self.days_of_month


class LatestOnlyMixin(ConditionRunMixin):
    """
    仅最新一次任务运行
    """
    def should_run(self, ctx):
        if ctx['dag_run'] and ctx['dag_run'].external_trigger:
            self.log.info("Externally triggered DAG_Run: allowing execution to proceed.")
            return True
        now = pendulum.now(tz=timezone('UTC'))
        left_window = ctx['dag'].following_schedule(ctx['execution_date'])
        right_window = ctx['dag'].following_schedule(left_window)
        self.log.info(
            'Checking latest only with left_window: %s right_window: %s now: %s',
            left_window, right_window, now
        )
        return left_window < now <= right_window